import networkx as nx
import numpy as np
import torch
from typing import Dict, Any, List, Optional
from app.services.dhgam.tsge_service import TemporalSpatialGraphEmbedding
from app.schemas.dhgam.tsge_base_schema import (
    InputParams, OutputParams, AlgorithmResponse, TaskStatus,
    TSGEOutputParams, NodeEmbedding, LinkPrediction
)
from app.schemas.metrics_schema import LineChart, ScatterPlot
from app.utils.logger import logger

class TSGEService:
    """时空图嵌入服务"""
    
    async def run_tsge(
        self,
        input_params: Dict[str, Any],
        task_id: str
    ) -> AlgorithmResponse:
        """
        运行TSGE算法并返回结果
        
        参数:
            input_params: 算法输入参数
            task_id: 任务ID
            
        返回:
            AlgorithmResponse: 算法响应，包含节点嵌入和链接预测结果
        """
        try:
            # 提取算法参数
            tsge_params = input_params.get("tsge_params", {})
            graph_sequence_data = tsge_params.get("graph_sequence", [])
            embedding_dim = tsge_params.get("embedding_dim", 128)
            learning_rate = tsge_params.get("learning_rate", 0.01)
            epochs = tsge_params.get("epochs", 100)
            
            # 直接使用原始图数据，不转换为NetworkX图
            # 初始化TSGE模型
            tsge = TemporalSpatialGraphEmbedding(
                embedding_dim=embedding_dim,
                learning_rate=learning_rate,
                epochs=epochs
            )
            
            # 训练模型，直接传入原始图序列数据
            tsge.fit(graph_sequence_data)
            
            # 提取节点嵌入
            node_embeddings = []
            for node, embedding in tsge.node_embeddings.items():
                # 从原始图数据中查找节点类型
                node_type = None
                for graph in graph_sequence_data:
                    for node_data in graph['nodes']:
                        if str(node_data['id']) == str(node):
                            node_type = str(node_data['type'])
                            break
                    if node_type:
                        break
                
                node_embeddings.append(NodeEmbedding(
                    node_id=str(node),
                    embedding=embedding.tolist(),
                    node_type=node_type
                ))
            
            # 生成链接预测
            link_predictions = []
            if graph_sequence_data:
                last_graph = graph_sequence_data[-1]
                nodes = [str(node['id']) for node in last_graph['nodes']]
                current_time = len(graph_sequence_data) - 1
                
                # 随机选择节点对进行预测
                import random
                sample_size = min(100, len(nodes) * (len(nodes) - 1) // 2)
                node_pairs = []
                for _ in range(sample_size):
                    u = random.choice(nodes)
                    v = random.choice(nodes)
                    if u != v:
                        node_pairs.append((u, v))
                
                # 获取节点类型映射
                node_type_map = {str(n['id']): str(n['type']) for n in last_graph['nodes']}
                
                # 预测链接概率
                for u, v in node_pairs:
                    u_type = node_type_map.get(u)
                    v_type = node_type_map.get(v)
                    prob = tsge.predict_link(u, v, current_time, u_type, v_type)
                    link_predictions.append(LinkPrediction(
                        source=str(u),
                        target=str(v),
                        probability=float(prob),
                        time=float(current_time)
                    ))
            
            # 创建算法输出参数
            output_params = OutputParams(
                tsge_results=TSGEOutputParams(
                    node_embeddings=node_embeddings,
                    link_predictions=link_predictions,
                    algorithm="TSGE",
                    parameters={
                        "embedding_dim": embedding_dim,
                        "learning_rate": learning_rate,
                        "epochs": epochs,
                        "num_graphs": len(graph_sequence_data),
                        "total_nodes": len(tsge.node_embeddings)
                    }
                )
            )
            
            # 创建可视化指标
            metrics = []
            
            # 如果有足够的节点，创建嵌入可视化
            if len(node_embeddings) > 5:
                # 降维并创建散点图
                from sklearn.decomposition import PCA
                embeddings_array = np.array([ne.embedding for ne in node_embeddings])
                pca = PCA(n_components=2)
                reduced_embeddings = pca.fit_transform(embeddings_array)
                
                scatter_data = []
                for i, ne in enumerate(node_embeddings):
                    scatter_data.append({
                        "x": float(reduced_embeddings[i, 0]),
                        "y": float(reduced_embeddings[i, 1]),
                        "id": ne.node_id,
                        "type": ne.node_type or "unknown"
                    })
                
                metrics.append(ScatterPlot(
                    title="节点嵌入降维可视化",
                    description="使用PCA降维的节点嵌入",
                    xAxis="PC1",
                    yAxis="PC2",
                    data=scatter_data
                ))
            
            # 创建响应对象
            response = AlgorithmResponse(
                task_id=task_id,
                task_status=TaskStatus.COMPLETED,
                task_progress=100,
                output_params=output_params,
                metrics=metrics
            )
            
            return response
            
        except Exception as e:
            logger.error(f"TSGE算法执行失败: {str(e)}")
            return AlgorithmResponse(
                task_id=task_id,
                task_status=TaskStatus.FAILED,
                error_message=f"TSGE算法执行失败: {str(e)}",
                task_progress=0,
                output_params=OutputParams()
            ) 